package suppressor

import (
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/cache"
	"github.com/stackrox/rox/generated/storage"
)

type cveSuppressorImpl struct {
	cache cache.VulnReqCache
}

func (m *cveSuppressorImpl) EnrichImageWithSuppressedCVEs(image *storage.Image) {
	vulnsMap := m.cache.GetVulnsWithState(image.GetName().GetRegistry(), image.GetName().GetRemote(), image.GetName().GetTag())
	if len(vulnsMap) == 0 {
		return
	}

	for _, comp := range image.GetScan().GetComponents() {
		for _, vuln := range comp.GetVulns() {
			vuln.State = vulnsMap[vuln.GetCve()]
		}
	}
}

func (m *cveSuppressorImpl) EnrichImageV2WithSuppressedCVEs(image *storage.ImageV2) {
	vulnsMap := m.cache.GetVulnsWithState(image.GetName().GetRegistry(), image.GetName().GetRemote(), image.GetName().GetTag())
	if len(vulnsMap) == 0 {
		return
	}

	for _, comp := range image.GetScan().GetComponents() {
		for _, vuln := range comp.GetVulns() {
			vuln.State = vulnsMap[vuln.GetCve()]
		}
	}
}
